import pandas as pd

from typing import List
from torch import Tensor
from sentence_transformers import util
from typing import Tuple, Dict

class PsychologyMatcher:
    def __init__(self, model, excel_path: str):
        """
        初始化类
        :param model:
        :param excel_path:
        """
        # 初始化模型
        self.model = model
        # 读取 Excel
        self.df = pd.read_excel(excel_path)
        # 把空值替换成空字符串，防止向量化报错
        self.df.fillna("", inplace=True)
        # 获取不重复的类别
        self.unique_categories = self.df['类别'].unique().tolist()
        # 加载数据，提取所有类别，并预先计算其向量嵌入，便于后续快速匹配。
        self.category_embeddings = self.model.encode(self.unique_categories, convert_to_tensor=True)

    def match_category(self, user_text: str, threshold: float = 0.4) -> List[Tuple[str, float]]:
        """
        将用户的输入与所有“类别”进行匹配，返回超过某个阈值的类别及其相似度分数。
        若没有类别超过阈值，则返回匹配度最高的一个类别。
        :param user_text:
        :param threshold:
        :return:
        """
        # 将用户输入转化成向量
        user_embedding = self.model.encode(user_text, convert_to_tensor=True)
        # 与所有类别对比
        scores: Tensor = util.cos_sim(user_embedding, self.category_embeddings)[0]
        # 找到分数超过某个阈值的下标
        matched = [
            (self.unique_categories[i], score.item())
            for i, score in enumerate(scores)
            if score.item() >= threshold
        ]
        # 如果没有匹配到任何一个超过阈值的类别，则选择匹配度最高的一个
        if not matched:
            max_index = scores.argmax().item()
            matched = [(self.unique_categories[max_index], scores[max_index].item())]
        # 可选：按相似度从高到低排序
        matched.sort(key=lambda x: x[1], reverse=True)

        # 输出格式是一个列表包含元组(匹配到的类别, 相似度分数)。
        return matched

    def match_issue_in_category(self, category: str, user_text: str, score_threshold: float = 0.7) -> List:
        """
        基于一级分类结果，进一步在对应类别下找出最相似的“心理问题+行为表现”，并返回其对应的心理问题、原因、干预方法等。
        :param category:
        :param user_text:
        :return:
        """
        # 筛选出该类别下的所有问题
        sub_df = self.df[self.df['类别'] == category]
        # 拼接“心理问题 + 行为表现”作为语义匹配文本
        behavior_embeddings = self.model.encode((sub_df['心理问题'] + "：" + sub_df['表现']).tolist(), convert_to_tensor=True)
        user_embedding = self.model.encode(user_text, convert_to_tensor=True)
        # 计算相似度分数
        scores = util.cos_sim(user_embedding, behavior_embeddings)[0]
        # 找出所有分数超过阈值的项
        matched_indices = (scores >= score_threshold).nonzero(as_tuple=True)[0]

        results = []

        if len(matched_indices) > 0:
            # 返回所有超过阈值的匹配项
            for idx in matched_indices.tolist():
                row = sub_df.iloc[idx]
                if row['表现'] is None:
                    row['表现'] = ''
                results.append({
                    "匹配心理问题": row['心理问题'],
                    "表现": row['表现'],
                    "产生原因": row['原因'],
                    "辅导方法": row['辅导方法'],
                    "问卷": row['问卷'],
                    # "匹配分数": scores[idx].item()
                })
        else:
            # 无结果超过阈值，返回最高分项
            max_idx = scores.argmax().item()
            row = sub_df.iloc[max_idx]
            if row['表现'] is None:
                row['表现'] = ''
            results.append({
                "匹配心理问题": row['心理问题'],
                "表现": row['表现'],
                "产生原因": row['原因'],
                "辅导方法": row['辅导方法'],
                "问卷": row['问卷'],
                # "匹配分数": scores[idx].item()
            })


        return results
